package com.mzy.linear_struct;

import java.util.LinkedList;
import java.util.List;
import java.util.Queue;
import java.util.function.Consumer;

/**
 * @Author: codermzy
 * @Date: 2024/04/18/15:59
 * @Description:
 */
public class AVLTree {

    protected AVLTreeNode root;

    protected void checkBalance(AVLTreeNode newNode) {

        AVLTreeNode current = newNode.parent;
        while (current != null) {
            if (!current.isBalanced()) {
                this.reBalance(current);
                break;
            }
            current = current.parent;
        }

    }

    public void reBalance(AVLTreeNode root) {
        AVLTreeNode pivot = root.higherChild();
        AVLTreeNode current = pivot.higherChild();

        if (pivot.isLeft()) { // L
            if (current.isLeft()) { // L L
                root.rightRotation();// 左左情况，进行右旋
            } else { // L R
                // 左右情况：选左旋，再右旋
                pivot.leftRotation();
                root.rightRotation();
            }
        } else { // R
            if (current.isLeft()) { // R L
                // 右左情况：选右旋，再左旋
                pivot.rightRotation();
                root.leftRotation();
            } else { // R R
                // 右右情况，进行左旋
                root.leftRotation();
            }
        }

    }


    private class AVLTreeNode {
        private Integer value;
        private AVLTreeNode left;
        private AVLTreeNode right;

        private AVLTreeNode parent;


        @Override
        public String toString() {
            return "AVLTreeNode{" +
                    "value=" + value +
                    '}';
        }

        public AVLTreeNode(Integer value) {
            this.value = value;
            this.left = null;
            this.right = null;
            this.parent = null;
        }

        public boolean isLeft() {
            return this.parent != null
                    && this.parent.left != null
                    && this.parent.left.equals(this);
        }

        public boolean isRight() {
            return this.parent != null
                    && this.parent.right != null
                    && this.parent.right.equals(this);
        }

        public int getHeight() {
            int leftHeight = this.left != null ? this.left.getHeight() : 0;
            int rightHeight = this.right != null ? this.right.getHeight() : 0;
            return Math.max(leftHeight, rightHeight) + 1;
        }

        // 获取权重因子（左边高度-右边高度）
        public int getBalanceFactor() {
            int leftHeight = this.left != null ? this.left.getHeight() : 0;
            int rightHeight = this.right != null ? this.right.getHeight() : 0;
            return leftHeight - rightHeight;
        }

        public boolean isBalanced() {
            int factor = this.getBalanceFactor();

            return factor >= -1 && factor <= 1;
        }


        public AVLTreeNode higherChild() {
            int leftHeight = this.left != null ? this.left.getHeight() : 0;
            int rightHeight = this.right != null ? this.right.getHeight() : 0;

            if (leftHeight > rightHeight) return this.left;
            if (rightHeight > leftHeight) return this.right;

            return this.isLeft() ? this.left : this.right;
        }

        // 左左情况：旋转操作：右旋转
        public AVLTreeNode rightRotation() {
            // root、pivot
            AVLTreeNode curRoot = this;
            AVLTreeNode pivot = this.left;
            boolean isLeft = this.isLeft();
            boolean isRight = this.isRight();
            // 1. 处理 pivot 节点

            // 2. 处理 pivot.right 节点

            // 3. 处理 this（root） 节点

            // 4. 挂载 pivot 节点


            // 1. 将 pivot 的父节点指向 root 的父节点
            pivot.parent = curRoot.parent;
            // 2.  root 节点的左子节点指向 pivot 的右子节点，pivot.right 有可能为空，但也没事
            curRoot.left = pivot.right;
            // 3. 将 pivot 的父节点属性改为 root，但因为 pivot.right 可能为空，所以要进行判空
            if (pivot.right != null) {
                pivot.right.parent = curRoot;
            }
            // 4. 把 pivot 的右子节点指向 root
            pivot.right = curRoot;
            // 5. root 的 parent属性 指向 pivot
            curRoot.parent = pivot;
            // 6. 最后一步：pivot应该挂在哪里
            // 情况一：pivot.parent 为空，那么它就是根节点，直接将 root = pivot
            if (pivot.parent == null) {
                root = pivot;
            } else if (isLeft) { // 情况三：pivot 为父节点的左子节点，那么就应该为：pivot.parent.left = pivot
                pivot.parent.left = pivot;
            } else if (isRight) {            // 情况二：pivot 为父节点的右子节点，那么就应该为：pivot.parent.right = pivot
                pivot.parent.right = pivot;
            }
            return pivot;
        }


        // 右右情况：旋转操作，左旋转
        public AVLTreeNode leftRotation() {
            // 获取 pivot、root
            AVLTreeNode pivot = this.right;
            AVLTreeNode curRoot = this;
            boolean isLeft = this.isLeft();
            boolean isRight = this.isRight();

            // 1. 处理 pivot
            pivot.parent = curRoot.parent;

            // 2. 处理 pivot.left
            curRoot.right = pivot.left;
            if (pivot.left != null) {
                pivot.left.parent = curRoot;
            }

            // 3. 处理 root
            pivot.left = curRoot;
            curRoot.parent = pivot;

            // 4. 处理挂载 pivot
            if (pivot.parent == null) {
                // 表示是根节点
                root = pivot;
            } else if (isLeft) {
                pivot.parent.left = pivot;
            } else if (isRight) {
                pivot.parent.right = pivot;
            }

            return pivot;
        }
    }


    public AVLTreeNode insert(Integer value) {

        AVLTreeNode newNode = this.createAVLTreeNode(value);

        // 判断是否是第一次插入
        if (root == null) {
            root = newNode;
        } else {
            insertNode(root, newNode);
        }
        checkBalance(newNode);
        return newNode;
    }


    protected AVLTreeNode createAVLTreeNode(Integer value) {
        return new AVLTreeNode(value);
    }

    public void preOrderTraverse(Consumer<Integer> consumer) {
        preOrderTraverseNode(root, consumer);
    }

    private void preOrderTraverseNode(AVLTreeNode root, Consumer<Integer> consumer) {
        if (root == null) {
            return;
        }
        consumer.accept(root.value);
        preOrderTraverseNode(root.left, consumer);
        preOrderTraverseNode(root.right, consumer);
    }

    public void inOrderTraverse() {
        inOrderTraverseNode(root);
    }

    private void inOrderTraverseNode(AVLTreeNode root) {
        if (root == null) {
            return;
        }
        inOrderTraverseNode(root.left);
        System.out.println(root);
        inOrderTraverseNode(root.right);
    }


    public void postOrderTraverse() {
        postOrderTraverseNode(root);
    }

    private void postOrderTraverseNode(AVLTreeNode root) {
        if (root == null) {
            return;
        }

        postOrderTraverseNode(root.left);
        postOrderTraverseNode(root.right);
        System.out.println(root);
    }


    private void insertNode(AVLTreeNode root, AVLTreeNode newNode) {


        if (root.value > newNode.value) {

            if (root.left == null) {
                root.left = newNode;
                newNode.parent = root;
            } else {
                insertNode(root.left, newNode);
            }
        } else if (root.value < newNode.value) {
            if (root.right == null) {
                root.right = newNode;
                newNode.parent = root;
            } else {
                insertNode(root.right, newNode);
            }
        }


    }


    public List<Integer> levelOrderTraverse() {

        if (root == null) {
            return null;
        }
        Queue<AVLTreeNode> queue = new LinkedList<>();
        LinkedList<Integer> list = new LinkedList<>();
        queue.offer(root);

        while (!queue.isEmpty()) {
            AVLTreeNode node = queue.poll();
            // System.out.println(node);
            list.add(node.value);
            if (node.left != null) {
                queue.offer(node.left);
            }
            if (node.right != null) {
                queue.offer(node.right);
            }
        }
        return list;
    }


    public int getMax() {
        AVLTreeNode current = root;
        int max = 0;
        while (current != null && current.right != null) {
            AVLTreeNode rightNode = current.right;
            max = rightNode.value;
            current = current.right;
        }
        return max;
    }

    public int getMin() {
        AVLTreeNode current = root;
        int min = 0;
        while (current != null && current.left != null) {
            AVLTreeNode leftNode = current.left;
            min = leftNode.value;
            current = current.left;
        }
        return min;
    }


    public boolean search(Integer value) {
        AVLTreeNode node = searchNode(value);
        return node != null;
    }

    private AVLTreeNode searchNode(Integer value) {
        AVLTreeNode current = root;
        AVLTreeNode parent = null;
        while (current != null) {
            if (value.equals(current.value)) return current;

            parent = current;
            if (value > current.value) current = current.right;
            if (value < current.value) current = current.left;

            current.parent = parent;
        }
        return null;
    }

    public boolean remove(Integer value) {

        AVLTreeNode current = searchNode(value);
        if (current == null) {
            return false;
        }
        // 当前 current 是要删除的元素
        // 1. 如果是叶子节点直接进行删除
        if (current.left == null && current.right == null) {
            // 判断当前节点在左边还是右边
            if (current.isLeft()) {
                current.parent.left = null;
            } else if (current.isRight()) {
                current.parent.right = null;
            } else { // 根节点的情况
                root = null;
            }
        }
        // 2.1 如果不是叶子节点且只有左子节点
        else if (current.right == null) {
            current.left.parent = current.parent;
            // 判断是否是根节点
            if (current == root) {
                root = current.left;
            } else if (current.isLeft()) {
                current.parent.left = current.left;
            } else {
                current.parent.right = current.left;
            }
        }
        // 2.2 如果不是叶子节点且只有右子节点
        else if (current.left == null) {
            current.right.parent = current.parent;
            if (current == root) {
                root = current.right;
            } else if (current.isRight()) {
                current.parent.right = current.right;
            } else {
                current.parent.left = current.right;
            }
        } else {      // 左右节点都有的情况
            // 前驱节点：比当前节点小一点点，但 是当前节点右子节点中最大的
            // 后继节点：比当前节点大一点点，但 是当前节点左子节点中最小的
            AVLTreeNode successor = getSuccessor(current);
            current.value = successor.value;
            current = successor;
            // if(current.isLeft()) {
            //     current.parent.left = successor;
            // }else {
            //     current.parent.right = successor;
            // }
            // 判断是否是根节点
            // if (current == root) {
            //     root = successor;
            // } else if (current.isLeft()) { // 不是根节点
            //     current.parent.left = successor;
            // } else if (current.isRight()) {
            //     current.parent.right = successor;
            // }

        }
        checkBalance(current);
        return true;
    }

    private AVLTreeNode getSuccessor(AVLTreeNode delNode) {
        // 拿到右边的节点
        AVLTreeNode current = delNode.right;
        AVLTreeNode successor = null;
        // 获取右边最小的节点
        while (current != null) {
            successor = current;
            current = current.left;
            if (current != null) {
                current.parent = successor;
            }
        }
        System.out.println("删除节点：" + delNode.value + "   后继节点：" + successor);


        if (successor != delNode.right) {
            assert successor != null;
            successor.parent.left = successor.right;
            // successor.right = delNode.right;
            if (successor.right != null) {
                successor.right.parent = successor.parent;
            }
        } else {
            assert successor != null;
            delNode.right = successor.right;
            if (successor.right != null) {
                successor.right.parent = delNode;
            }
        }

        // assert successor != null;
        // successor.left = delNode.left;
        return successor;
    }
}
